/*
1519.子树中标签相同的节点数
https://leetcode.cn/problems/number-of-nodes-in-the-sub-tree-with-the-same-label/
中等 2024/11/19
*/

#include <vector>
#include <unordered_map>
#include <string>
#include <iostream>
using namespace std;
class Solution {
public:
    vector<int> countSubTrees(int n, vector<vector<int>>& edges, string labels) {
        vector<int> ans(n, 0);
        vector<vector<int>> tree(n);
        
        // 构建邻接表
        for (const auto& edge : edges) {
            int u = edge[0], v = edge[1];
            tree[u].push_back(v);
            tree[v].push_back(u);
        }

        // DFS辅助函数
        function<unordered_map<char, int>(int, int)> dfs = [&](int node, int parent) -> unordered_map<char, int> {
            unordered_map<char, int> labelCount;
            labelCount[labels[node]] = 1; // 当前节点标签计数

            // 遍历所有子节点
            for (int neighbor : tree[node]) {
                if (neighbor == parent) continue; // 避免回到父节点

                // 获取子树的标签计数
                auto childLabelCount = dfs(neighbor, node);

                // 将子树的标签计数合并到当前节点的标签计数
                for (const auto& entry : childLabelCount) {
                    labelCount[entry.first] += entry.second;
                }
            }

            // 最后计算当前节点的标签在其子树中的出现次数
            ans[node] = labelCount[labels[node]];
            
            return labelCount;
        };

        // 从根节点0开始DFS
        dfs(0, -1);
        return ans;
    }
};
// 测试用例
int main() {
    Solution solution;
    int n1 = 7;
    vector<vector<int>> edges1 = {{0, 1}, {0, 2}, {1, 4}, {1, 5}, {2, 3}, {2, 6}};
    string labels1 = "abaedcd";
    vector<int> result1 = solution.countSubTrees(n1, edges1, labels1);
    for (int i : result1) cout << i << " ";  // 输出：2 1 1 1 1 1 1
    cout << endl;

    int n2 = 4;
    vector<vector<int>> edges2 = {{0, 1}, {1, 2}, {0, 3}};
    string labels2 = "bbbb";
    vector<int> result2 = solution.countSubTrees(n2, edges2, labels2);
    for (int i : result2) cout << i << " ";  // 输出：4 2 1 1
    cout << endl;

    return 0;
}